package com.ics.atable.chat.serializer;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import com.ics.atable.chat.model.vo.node.NodeResponse;
import com.ics.atable.chat.model.vo.node.IntentData;
import com.ics.atable.chat.model.vo.node.StructTokenData;
import com.ics.atable.chat.model.vo.node.RetrievalData;
import com.ics.atable.chat.model.vo.node.stream.ObjectStreamResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Centralized codec for assistant content serialization/deserialization.
 */
@Slf4j
@Component
public class ChatContentCodec {

    private static final ObjectMapper OBJECT_MAPPER;

    static {
        OBJECT_MAPPER = new ObjectMapper();
        OBJECT_MAPPER.registerModule(new JavaTimeModule());
        OBJECT_MAPPER.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
    }

    /**
     * Serialize a list of ObjectStreamResponse (collected during streaming) into a JSON string
     * suitable for storing in DB.
     */
    public String serializeObjectStreamResponses(List<ObjectStreamResponse> responses) {
        try {
            if (responses == null) {
                return "[]";
            }
            // Flatten to list of NodeResponse lists for compact storage
            List<List<NodeResponse>> payload = new ArrayList<>();
            for (ObjectStreamResponse response : responses) {
                payload.add(response.getDataValue());
            }
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (Exception e) {
            log.error("Failed to serialize ObjectStreamResponses", e);
            // Fallback to empty array to avoid null persistence
            return "[]";
        }
    }

    /**
     * Deserialize stored JSON content back to ObjectStreamResponse.
     * Supports both array form (list of chunks) and single object/list form.
     */
    public ObjectStreamResponse deserializeObjectStreamResponse(String content) {
        if (content == null || content.isEmpty()) {
            return new ObjectStreamResponse(new ArrayList<>());
        }
        try {
            String trimmed = content.trim();
            if (trimmed.startsWith("[")) {
                // Either [[NodeResponse...], [NodeResponse...]] or [NodeResponse...]
                // Try parse as list of list first
                try {
                    List<List<NodeResponse>> responseLists = OBJECT_MAPPER.readValue(trimmed,
                            new TypeReference<List<List<NodeResponse>>>() {});
                    List<NodeResponse> all = new ArrayList<>();
                    for (List<NodeResponse> list : responseLists) {
                        for (NodeResponse nodeResponse : list) {
                            convertNodeDataType(nodeResponse);
                            all.add(nodeResponse);
                        }
                    }
                    return new ObjectStreamResponse(all);
                } catch (Exception ignore) {
                    // Fallback: parse as single list
                    List<NodeResponse> nodeResponses = OBJECT_MAPPER.readValue(trimmed,
                            new TypeReference<List<NodeResponse>>() {});
                    for (NodeResponse nodeResponse : nodeResponses) {
                        convertNodeDataType(nodeResponse);
                    }
                    return new ObjectStreamResponse(nodeResponses);
                }
            }
            if (trimmed.startsWith("{")) {
                // Some historical records may store a single list as JSON object; parse as list
                List<NodeResponse> nodeResponses = OBJECT_MAPPER.readValue(trimmed,
                        new TypeReference<List<NodeResponse>>() {});
                for (NodeResponse nodeResponse : nodeResponses) {
                    convertNodeDataType(nodeResponse);
                }
                return new ObjectStreamResponse(nodeResponses);
            }
            log.warn("Content is not valid JSON for ObjectStreamResponse: {}", content);
            return new ObjectStreamResponse(new ArrayList<>());
        } catch (Exception e) {
            log.error("Failed to deserialize ObjectStreamResponse, content: {}", content, e);
            return new ObjectStreamResponse(new ArrayList<>());
        }
    }

    private static final Map<String, Class<?>> NODE_TO_DATA_CLASS = new HashMap<>() {{
        put("意图识别", IntentData.class);
        put("结构化提取", StructTokenData.class);
        put("指标检索", RetrievalData.class);
    }};

    private void convertNodeDataType(NodeResponse nodeResponse) {
        if (nodeResponse == null) {
            return;
        }
        Object raw = nodeResponse.getData();
        if (raw == null) {
            return;
        }
        Class<?> target = NODE_TO_DATA_CLASS.get(nodeResponse.getNode());
        if (target == null) {
            return;
        }
        try {
            Object converted = OBJECT_MAPPER.convertValue(raw, target);
            nodeResponse.setData(converted);
        } catch (IllegalArgumentException e) {
            log.warn("Failed to convert node '{}' data to {}", nodeResponse.getNode(), target.getSimpleName());
        }
    }
}


